from class_task import *

#每个Stage为一个Node
class Node():
    def __init__(self,stage,dag) -> None:
        self.id = stage[0] #注意，是从1开始的
        self.unit = stage[1]
        self.time = stage[2]
        self.pre = list(stage[3]) #前置stage
        self.sub = [] #后置stage

        #如果定义了要访问的页
        if len(stage) > 4:
            self.die = stage[4][0]
            self.page = stage[4][1] #访问的Page
        else:
            self.die = -1
            self.page = -1

        #依赖度
        self.depen = -1 #初始化为-1表示尚未被计算

        #下列结构体是为了调度方便，查看是否调度完成了
        #指向整张图
        self.DAG = dag
        self.finish_time = -1 #完成时间

    def fresh(self):
        self.finish_time = -1 #完成时间
        
#每个Task为一个DAG
class DAG():
    def __init__(self,task) -> None:
        self.head = None #整个DAG的初始节点
        self.name = task.task_name
        self.node_list = [] #所有Node按id为索引存储在list中，便于按照id寻找node

        #将每个stage初始化为一个Node
        for stage in task.task:
            node = Node(stage,self)
            self.node_list.append(node)
            if node.id == 1:
                self.head = node

        #添加后置任务
        for node in self.node_list:
            #将这个node自己的id添加到前置任务节点的sub数组里面
            #1除外，否则就添加到最后一个node里面去了
            if node.id == 1:
                continue
            for pre_stage in node.pre:
                self.node_list[pre_stage-1].sub.append(node.id)

        #得到所有任务的依赖度
        self.get_all_depen()

    def fresh(self):
        for node in self.node_list:
            node.fresh()

    #查询是否具有依赖关系
    def is_parent(self,parent_id,child_id):
        #如果都找到初始节点了都没找到那就没有依赖关系
        if child_id == 0:
            return False

        child_node = self.node_list[child_id-1]
        for pre_id in child_node.pre:
            #找到与查询节点id相同的节点了
            if pre_id == parent_id:
                return True
            #没有找到的话继续找父节点的父节点
            else:
                return self.is_parent(parent_id,pre_id)

    #计算依赖度
    def get_all_depen(self):
        #分别计算每个子节点的依赖度
        def get_depen(node_id):
            node = self.node_list[node_id-1]
            #如果之前计算过了就直接获取
            if node.depen != -1:
                return node.depen

            # 木有计算过的话那就算一下并且存起来
            # 木有子节点的话那依赖度就是0
            if len(node.sub)==0:
                node.depen = 0
                return 0

            # 有子节点的话依赖度是所有子节点的（运行时间+依赖度）的最大值
            depen = 0
            for child_id in node.sub:
                child_run_time = self.node_list[child_id-1].time
                depen = max(depen,child_run_time+get_depen(child_id))
            node.depen = depen
            return depen

        for node in self.node_list:
            get_depen(node.id)

    #用于打印
    def __str__(self) -> str:
        return "DAG with "+str(len(self.node_list))+" Node"

if __name__=="__main__":
    read1 = ["读1",(1,"nfc1",100,(0,)),(2,"nfc1",100,(1,)),(3,"host",20,(2,)),(4,"host",5,(3,))]
    read2 = ["读2",(1,"nfc2",100,(0,)),(2,"nfc2",100,(1,)),(3,"host",20,(2,)),(4,"host",5,(3,))]
    write1 = ["写",(1,"nfc1",100,(0,)),(2,"nfc1",100,(1,)),(3,"host",20,(2,)),(4,"nfc1",600,(3,)),(5,"host",5,(4,))]
    text1 = ["文本",(1,"host",20,(0,)),(2,"nfc1",200,(1,)),(3,"nfc1",200,(2,)),(4,"matcher",150,(3,)),(5,"nfc2",200,(2,)),(6,"matcher",150,(5,)),(7,"host",20,(4,6,))]

    task = Task(text1)

    dag = DAG(task)

    for node in dag.node_list:
        print(node.id,node.pre,node.sub,node.depen)

    print(dag.is_parent(4,6))